#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#include <error.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <unistd.h>
#include <errno.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <poll.h>

#include <openssl/aes.h>

//gcc -Wall aes_server.c -o aes_server -lssl -l crypto

struct aes_key_t{
    AES_KEY encrypt_key;
    AES_KEY decrypt_key;
};

struct aes_key_t* aes_encrypt_init(void);
void aes_encrypt(unsigned char *in,unsigned char *out,size_t length,const struct aes_key_t* key);
void aes_decrypt(unsigned char *in,unsigned char *out,size_t length,const struct aes_key_t* key);
void aes_encrypt_deinit(struct aes_key_t* key);

#define SERVER_ADDR "127.0.0.1"
#define SERVER_PORT 5050

int main(int argc,char *argv[])
{
    int ret = -1;
    int sock_fd,conn_fd;
    struct sockaddr_in server_addr;
    unsigned char buf[sizeof(struct in_addr)];
    uint8_t recv_buf[64];
    ssize_t len;
    struct pollfd fds[2];
    char input_buf[12];
    struct aes_key_t *aes_key;
    unsigned char *out_buf;
    size_t encrypt_len;

    printf("run server:\r\n");

    aes_key = aes_encrypt_init();
    if(!aes_key)
    {
        printf("aes_encrypt_init error:%d..\r\n",ret);
        return ret;
    }

    sock_fd = socket(AF_INET,SOCK_STREAM,0);
    if(sock_fd < 0)
    {
        perror("socket ");
        return ret;
    }

    int opt = 1;
    setsockopt(sock_fd,SOL_SOCKET,SO_REUSEADDR,&opt,sizeof(opt));

    ret = inet_pton(AF_INET,SERVER_ADDR,buf);
    if(ret != 1)
    {
        printf("inet_pton error..\r\n");
        goto close_sock_fd;
    }
    server_addr.sin_addr.s_addr = ((struct in_addr *)buf)->s_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(SERVER_PORT);
    ret = bind(sock_fd,(struct sockaddr *)&server_addr,sizeof(server_addr));
    if(ret < 0)
    {
        perror("bind ");
        goto close_sock_fd;
    }
    ret = listen(sock_fd,250);
    if(ret < 0)
    {
        perror("listen ");
        goto close_sock_fd;
    }

    printf("wait client connect..\r\n");

    conn_fd = accept(sock_fd,NULL,NULL);
    if(conn_fd < 0)
    {
        perror("accept ");
        goto close_sock_fd;
    }

    printf("client connect success..\r\n");

    fds[0].fd = 0;
    fds[0].events = POLLIN;
    fds[1].fd = conn_fd;
    fds[1].events = POLLIN;
    for(;;)
    {
        ret = poll(fds,sizeof(fds) / sizeof(*fds),-1);
        if(ret < 0)
        {
            perror("poll ");
            goto close_sock_fd;
        }
        else if(!ret)
        {
            /*timeout*/
        }
        else
        {
            for(int i = 0;i < sizeof(fds) / sizeof(*fds);i++)
            {
                if(fds[i].revents & fds[i].events)
                {
                    if(fds[i].fd == 0)
                    {
                        ret = read(0,input_buf,sizeof(input_buf));
                        if(ret < 0)
                        {
                            if((ret != EINTR) && (ret != EAGAIN) && (ret != EWOULDBLOCK))
                            {
                                perror("read ");
                                goto close_conn_fd;
                            }
                        }
                        else if(!ret)
                        {

                        }
                        else
                        {
                            encrypt_len = (ret / 16 + 1) * 16;
                            out_buf = (unsigned char *)calloc(1,encrypt_len);
                            if(!out_buf)
                            {
                                printf("malloc failed..\r\n");
                            }
                            else
                            {
                                aes_encrypt((unsigned char *)input_buf,out_buf,ret,aes_key);
                                send(conn_fd,out_buf,encrypt_len,0);

                                free(out_buf);
                            }
                        }
                    }
                    else if(fds[i].fd == conn_fd)
                    {
                        memset(recv_buf,0,sizeof(recv_buf));
                        len = recv(conn_fd,recv_buf,sizeof(recv_buf),0);
                        if(len < 0)
                        {
                            if((len != EINTR) && (len != EAGAIN) && (len != EWOULDBLOCK))
                            {
                                perror("recv ");
                                goto close_conn_fd;
                            }
                            continue;
                        }
                        else if(!len)
                        {
                            printf("client shutdown..\r\n");
                            goto close_conn_fd;
                        }
                        else
                        {
                            out_buf = (unsigned char *)calloc(1,len);
                            if(!out_buf)
                            {
                                printf("malloc failed..\r\n");
                            }
                            else
                            {
                                aes_decrypt((unsigned char *)recv_buf,out_buf,len,aes_key);
                                printf("recv:%s\r\n",out_buf);

                                free(out_buf);
                            }
                        }
                    }
                }
            }
        }
    }

    ret = 0;
close_sock_fd:
    close(sock_fd);
close_conn_fd:
    close(conn_fd);

    return ret;
}

struct aes_key_t* aes_encrypt_init(void)
{
    int ret = -1;
    const unsigned char usr_key[] = {   0x00,0x11,0x22,0x33,0x44,0x55,0x66,0x77,
                                        0x88,0x99,0xaa,0xbb,0xcc,0xdd,0xee,0xff};

    struct aes_key_t *key;

    key = (struct aes_key_t *)calloc(1,sizeof(struct aes_key_t));
    if(!key)
    {
        printf("%s,malloc failed..\r\n",__FUNCTION__);
        return NULL;
    }

    ret = AES_set_encrypt_key(usr_key,sizeof(usr_key) * 8,&key->encrypt_key);
    if(ret < 0)
    {
        printf("AES_set_encrypt_key error:%d..\r\n",ret);
        return NULL;
    }
    ret = AES_set_decrypt_key(usr_key,sizeof(usr_key) * 8,&key->decrypt_key);
    if(ret < 0)
    {
        printf("AES_set_decrypt_key error:%d..\r\n",ret);
        return NULL;
    }

    return key;
}

void aes_encrypt(unsigned char *in,unsigned char *out,size_t length,const struct aes_key_t* key)
{
    unsigned char ivec[16];

    memset(ivec,0,sizeof(ivec));
    AES_cbc_encrypt(in,out,length,&key->encrypt_key,ivec,AES_ENCRYPT);
}

void aes_decrypt(unsigned char *in,unsigned char *out,size_t length,const struct aes_key_t* key)
{
    unsigned char ivec[16];

    memset(ivec,0,sizeof(ivec));
    AES_cbc_encrypt(in,out,length,&key->decrypt_key,ivec,AES_DECRYPT);
}

void aes_encrypt_deinit(struct aes_key_t* key)
{
    free(key);
}
